import PyPDF2
from langchain.text_splitter import RecursiveCharacterTextSplitter
import numpy as np
from transformers import AutoTokenizer, AutoModel
import os
import torch
import faiss
import numpy as np

def read_pdf(file_path):
    """
    使用PyPDF2提取目标PDF文件的内容。
    使用RecursiveCharacterTextSplitter分割PDF文本。

    该函数将传入的PDF文本分割成多个较小的文本块，以便于后续处理。
    分割时尽量保持句子或段落的完整性，并允许相邻文本块之间有一定的重叠。

    参数:
    pdftext (str): 需要分割的PDF文本。

    返回:
    list: 分割后的文本块列表。
    """
    try:
        with open(file_path, 'rb') as file:
            pdf_reader = PyPDF2.PdfReader(file)
            text = ""
            for page_num in range(len(pdf_reader.pages)):
                page = pdf_reader.pages[page_num]
                text += page.extract_text()
            return text
    except FileNotFoundError:
        print("错误：文件未找到！")
    except Exception as e:
        print(f"错误：发生了一个未知错误：{e}")
    return None

def EmbadderStr(pdftext):
    """
    使用RecursiveCharacterTextSplitter分割PDF文本。

    该函数将传入的PDF文本分割成多个较小的文本块，以便于后续处理。
    分割时尽量保持句子或段落的完整性，并允许相邻文本块之间有一定的重叠。

    参数:
    pdftext (str): 需要分割的PDF文本。

    返回:
    list: 分割后的文本块列表。
    """
    try:
#RecursiveCharacterTextSplitter 类用于递归地按字符对文本进行切片，避免将完整的句子或段落切断。
#采用递归式分割，会优先按特定分隔符（如 \n\n、\n、 等）来分割文本，尽量保证分割出的块是完整的句子、段落等，若按分隔符分割后块仍超出 chunk_size，才会进一步按字符进行切割。
#分割单位不仅可以是句子，还能根据需求调整为段落等更大的语义单元，并且可以控制分割块的长度，在保证语义完整的同时，对分割结果的长度有一定约束。
#可以将段落作为分割的参考，但不局限于段落，当段落过长时会进一步分割，同时可以设置重叠部分以保留上下文信息。
#适用于需要对段落进行细化处理，同时又要保证分割块之间有一定上下文关联的场景，如文档的结构化分析、信息检索等。
# 创建文本分割器
#chunk_size     定义了每个切片块（chunk）的最大字符长度。当对文本进行切片时，RecursiveCharacterTextSplitter 会尝试将文本切分成不超过 chunk_size 长度的块。
#chunk_overlap  表示相邻两个切片块之间重叠的字符数量。设置重叠部分是为了确保在处理文本时，不会因为切片而丢失重要的上下文信息。
#length_function    是一个可调用对象（通常是函数），用于计算文本的长度。默认情况下，使用 Python 内置的 len 函数来计算字符数量。你也可以自定义这个函数，例如根据词数、字节数等来计算文本长度。
#separators     是一个列表，用于定义文本分割的分隔符。默认情况下，RecursiveCharacterTextSplitter 会根据句子的结束符（如句号、问号、感叹号等）进行分割。
#keep_separator     表示是否保留分隔符。默认为 False，表示不保留分隔符。设置为 True 时，分割后的文本块将包含分隔符。
        text_splitter = RecursiveCharacterTextSplitter(
            chunk_size = 200,       #切分后的每个文本块的字符数量最多为 200。但需要注意的是，为了保证句子或段落的完整性，实际切分的块长度可能会小于 chunk_size。
            chunk_overlap  = 20,    #相邻两个切片块会有 20 个字符是重复的。例如，第一个切片块的最后 20 个字符会和第二个切片块的前 20 个字符相同。
            length_function = len,
            separators=["。", "！", "？", "\n\n", "\n", " "],
            keep_separator=True,
            )
        pdfsentences=text_splitter.split_text(pdftext)
        return pdfsentences
    except Exception as e:
        print(f"错误：发生了一个未知错误：{e}")
    return None


# 按句子切片
#严格按照句子的边界（如句号、问号、感叹号等）进行分割，每个分割块就是一个完整的句子。
#常用于需要对单个句子进行分析的场景，如情感分析、词性标注等，聚焦于句子层面的语义理解。
#pdf_sentences = nltk.sent_tokenize(pdf_text)


def vectorizerKnowledgeStr(pdfsentences,tokenizerk,modelk):
    """
    将文本数据转换为向量表示。

    参数:
    pdfsentences (list): 包含文本块的列表，每个块将被转换为向量。
    tokenizerk: 用于文本分词的模型。
    modelk: 用于生成文本向量的预训练模型。

    返回:
    torch.Tensor: 包含所有文本块向量的张量，如果发生错误则返回None。
    """
    try:
        # 文本向量化
        vectors = []
        for chunk in pdfsentences:
            inputs = tokenizerk(chunk, return_tensors='pt')
            with torch.no_grad():
                outputs = modelk(**inputs)
            # 这里简单取 [CLS] 标记的输出作为句子向量
            #vector = outputs.last_hidden_state[:, 0, :].squeeze().numpy()    
            # 使用平均池化获取句子向量
            vector = torch.mean(outputs.last_hidden_state, dim=1).squeeze().numpy()
            vectors.append(vector)

        return torch.tensor(vectors)
    except Exception as e:
        print(f"错误：发生了一个未知错误：{e}")
    return None

def vectorDBIndex(knowledgevectors):
    """
    创建并返回一个基于给定向量集的 Faiss 索引。

    参数:
    knowledgevectors -- 一个二维数组，包含一系列知识向量。

    返回:
    如果成功创建索引，则返回一个 Faiss 索引对象；否则返回 None。
    """
    try:
        # 获取向量的维度
        dimension = knowledgevectors.shape[1]

        # 创建 Faiss 索引
        index = faiss.IndexFlatL2(dimension)  # 使用 L2 距离的扁平索引
        index.add(knowledgevectors)  # 将向量添加到索引中
        return index
    except Exception as e:
        print(f"错误：发生了一个未知错误：{e}")
    return None

def vectorizerQueryStr(queryText,tokenizerq,modelq):
    """
    将查询字符串转换为向量。

    参数:
    queryText: 查询字符串。
    tokenizerq: 用于分词的分词器对象。
    modelq: 用于编码向量的编码器对象。

    返回:
    一个包含查询向量的二维数组。
    """
    try:
        inputs = tokenizerq(queryText, return_tensors='pt')
        with torch.no_grad():
            outputs = modelq(**inputs)
        queryvector = torch.mean(outputs.last_hidden_state, dim=1).squeeze().numpy().reshape(1, -1)   
        return queryvector
    except Exception as e:
        print(f"错误：发生了一个未知错误：{e}")
    return None



if __name__ == '__main__':
    
    # 使用示例
    file_path = 'D:\VS_workspace\A002_AIPython学习\示例资料及练习\Modbus协议中文版(GB).pdf'
    if os.path.exists(file_path):
        pdf_text = read_pdf(file_path)
        print("已经成功从PDF读取到字符串，对PDF进行句子切片\r\n")

        pdf_sentences = EmbadderStr(pdf_text)

        # 打印切片结果
        print(f"PDF 句子切片结果[{len(pdf_sentences)}]-------------------------:\r\n")

        i=0
        #while i<len(pdf_sentences):print(f"句子 {i+1}: {pdf_sentences[i]}");i=i+1

        local_model_path = r"D:\bert_base_chinese_cache\models--bert-base-chinese"
        # 加载分词器和模型
        tokenizer = AutoTokenizer.from_pretrained(local_model_path)
        model = AutoModel.from_pretrained(local_model_path)

        knowledge_vectors=vectorizerKnowledgeStr(pdf_sentences,tokenizer,model)

        index = vectorDBIndex(knowledge_vectors) 

        # 执行相似性搜索
        query_text = "如何读取输入寄存器"

        query_vector = vectorizerQueryStr(query_text,tokenizer,model)

        k = 3    # 查找最相似的 2 个向量
        distances, indices = index.search(query_vector, k)

        # 输出搜索结果
        print("查询文本:", query_text)
        for i in range(k):
            print(f"第 {i + 1} 个最相似的文本块:")
            print(pdf_sentences[indices[0][i]])
            print(f"距离: {distances[0][i]}")
            print()

    print ("退出主程序")


"""
该程序
知识样本文件： file_path = 'D:\VS_workspace\A002_AIPython学习\示例资料及练习\Modbus协议中文版(GB).pdf'
提问： query_text = "如何读取输入寄存器"
使用了 PyPDF2 库来读取 PDF 文件；
使用了 RecursiveCharacterTextSplitter 库来将 PDF 文件拆分切片成 smaller chunks;
使用了 transformers 库中本地化的bert-base-chinese模型来计算知识文本及问题文本的嵌入向量；
使用了 faiss 库来加速计算，并使用 numpy 库来处理数组；

注意：
1、transformers 库在默认使用时是会向服务器请求模型的， 正常情况下是无法直接访问地址https://huggingface.co/的，所以可能无法正常运行。

BERT （ Bidirectional Encoder Representations from Transformers ） [ bert-base-chinese] ：谷歌开发，基于 Transformer 编码器，
    使用大规模语料库训练。它在语言理解任务上表现卓越，如文本分类、命名实体识别、问答系统等，通过对输入文本进行掩码处理来预测被掩盖的词，
    能很好地捕捉文本的上下文信息，但在文本生成任务上相对较弱。

运行结果：成功
2025-05-01 

"""    